#include "DoglegSolver.h"

namespace Solver
{
DoglegSolver::DoglegSolver() : BaseSolver(SolverType_E::Dogleg_E, "Dogleg"), alpha_(0), beta_(0), thres_(0), bestError_(1e300), currentError_(1e300)
{

}

DoglegSolver::~DoglegSolver()
{

}

bool DoglegSolver::init()
{
    //1.
    if (vecVariable_.empty())
    {
        return false;
    }

    //2.   
    bestX_.resize(fun_->getDim());
    y_.resize(fun_->getErrorDim());
    
    //2.
    auto &data = vecVariable_.front()->current->x;
    std::copy(data.begin(), data.end(), bestX_.data());
    bestError_    = calcError(bestX_);
    currentError_ = bestError_;

    //3.
    x_.resize(fun_->getDim());
    g_.resize(fun_->getDim());
    vDx_.resize(fun_->getDim());
    jacobi_.resize(fun_->getErrorDim(), fun_->getDim());
    std::copy(bestX_.data(), bestX_.data() + bestX_.size(), x_.data());
    
    //4.
    thres_ = fun_->getDim() * 10000.0;
    return true;
}


void DoglegSolver::solveByIter()
{
    //1.
    calcJacobi(x_);

    //2.
    calcG();

    //3.
    value_t v = g_.squaredNorm();
    if (std::isnan(v) || v < 1e-24)
    {
        condition_.stop.store(true, std::memory_order_relaxed);
        return;
    }

    //4.
    calcAlpha();

    //5.
    if(!solveVDX())
    {
        condition_.stop.store(true, std::memory_order_relaxed);
        return;
    }

    //6.
    modifyDx();

    //7.
    if (vDx_.squaredNorm() < 1e-16)
    {
        condition_.stop.store(true, std::memory_order_relaxed);
        return;
    }
}

void DoglegSolver::calcG()
{
    g_ = jacobi_.transpose() * y_;
}

void DoglegSolver::calcAlpha()
{
    alpha_ = g_.squaredNorm() / ((jacobi_ * g_).squaredNorm() + 1e-16);
}

void DoglegSolver::calcBeta()
{
    value_t c(0), a2(0), ba2(0), a;
    const value_t alpha2 = std::pow(alpha_, 2);
    Eigen::Index num = vecMax_.size();
    for (Eigen::Index i = 0; i != num; ++i)
    {
        a    = alpha2 * g_[i];
        a2  += std::pow(a, 2);
        c   += a * (vDx_[i] + a);
        ba2 += std::pow(vDx_[i] + a, 2);
    }

    //2.
    const value_t z = std::sqrt(std::pow(c, 2) + ba2 * (std::pow(thres_, 2) - a2));
    beta_ = (c < 0.0 ?  (std::pow(thres_, 2) - a2) / (z - c) : (z + c) / (ba2 + 1e-16));
}

void DoglegSolver::calcJacobi(const VectorXd_t &x)
{
    fun_->calcJacobiDiff(x, jacobi_);
}

bool DoglegSolver::solveVDX()
{
    //1.
    lineSolver_.solveLineEqu(jacobi_, y_, vDx_);

    //2.
    for (auto v : vDx_)
    {
        if (std::isnan(v))
        {
            return false;
        }
    }

    //3.
    return true;
}

void DoglegSolver::modifyDx()
{
    //1.
    std::transform(vDx_.begin(), vDx_.end(), vDx_.begin(), [](value_t x)->value_t
    {
        return -x;
    });

    //2.
    const value_t thres2 = thres_ * thres_;
    if (vDx_.squaredNorm() < thres2)
    {
        return;
    }

    //3.
    const value_t s = g_.squaredNorm();
    if (std::pow(alpha_, 4) * s > thres2)
    {
        value_t mul = thres_ / std::sqrt(s);
        if (alpha_ > 0)
        {
            mul = -mul;
        }

        std::transform(g_.begin(), g_.end(), vDx_.begin(), [&](value_t x)->value_t
        {
            return mul * x;
        });
    }

    //4.
    calcBeta();

    //5.
    value_t alpha2 = alpha_ * alpha_;
    std::transform(g_.begin(), g_.end(), vDx_.begin(), vDx_.begin(), [&](value_t g, value_t dx)->value_t
    {
		const value_t a = alpha2 * g;
		return beta_ * (dx + a) - a;
    });
}

value_t DoglegSolver::calcError(const VectorXd_t &x)
{
    fun_->calcErrorFun(x, y_);
    value_t v = y_.squaredNorm();
    if (std::isnan(v))
    {
        return 1e300;
    }
    else if(std::isinf(v))
    {
        return 1e300;
    }
    else
    {
        return v;
    }
}

}